import sys
from argparse import (
    RawTextHelpFormatter,
    ArgumentParser
)
from rich import print as rprint
from rich.console import Console
from utils import (
    list_models,
    multiline_input,
    handle_chat_commands,
    print_models,
    BANNER,
    init_env,
    SUPPORTED_APIS,
    PROMPT_SYMBOL,
    MULTILINE_END,
    ERROR_MESSAGES,
    SUCCESS_MESSAGES
)

console = Console()

def validate_api_and_model(api: str, model: str) -> bool:
    if api not in SUPPORTED_APIS:
        rprint(f'[bold red]{ERROR_MESSAGES["invalid_api"].format(api=api)}[/]')
        rprint(f'[yellow]Supported APIs: {", ".join(SUPPORTED_APIS)}[/]')
        return False
    available_models = list_models(api)
    if model not in available_models:
        rprint(f'[bold red]{ERROR_MESSAGES["invalid_model"].format(model=model)}[/]')
        rprint(f'[yellow]Available models for {api}: {", ".join(available_models)}[/]')
        return False
    return True

def run_command(args) -> None:
    if not validate_api_and_model(args.api, args.model):
        return
    rprint(f'[grey54]Model using: {args.api}.{args.model}[/grey54]')
    rprint('[dim]Type /help for available commands[/dim]')
    try:
        while True:
            try:
                content = input(PROMPT_SYMBOL).strip()
                if not content:
                    continue
                if content == MULTILINE_END:
                    content = multiline_input()
                result = handle_chat_commands(content, args.api, args.model)
                args.api, args.model = result
            except KeyboardInterrupt:
                rprint('\n[yellow]Use /q or /bye to exit[/yellow]')
                continue
            except EOFError:
                rprint('\n[bold green]<<<[/bold green] [bold magenta]Bye![/bold magenta]')
                break
    except Exception as e:
        rprint(f'[bold red]{ERROR_MESSAGES["unexpected_error"].format(error=e)}[/]')
        sys.exit(1)

def list_command(args) -> None:
    if args.api not in SUPPORTED_APIS:
        rprint(f'[bold red]{ERROR_MESSAGES["invalid_api"].format(api=args.api)}[/]')
        rprint(f'[yellow]Supported APIs: {", ".join(SUPPORTED_APIS)}[/]')
        return
    print_models(args.api)

def init_command(args) -> None:
    try:
        init_env()
        rprint(f'[bold green]{SUCCESS_MESSAGES["env_initialized"]}[/]')
    except Exception as e:
        rprint(f'[bold red]Failed to initialize environment: {e}[/]')

def create_parser() -> ArgumentParser:
    parser = ArgumentParser(
        formatter_class=RawTextHelpFormatter,
        description=BANNER,
        epilog="Examples:\n  shellgpt run ollama llama2\n  shellgpt ls openai\n  shellgpt init"
    )
    subparsers = parser.add_subparsers(
        metavar='COMMANDS',
        dest='command',
        help='Available commands'
    )
    # Run command
    run_parser = subparsers.add_parser(
        'run', 
        help='Start interactive chat session'
    )
    run_parser.add_argument(
        'api', 
        choices=SUPPORTED_APIS,
        help='API provider to use'
    )
    run_parser.add_argument(
        'model', 
        help='Model name to use'
    )
    run_parser.set_defaults(func=run_command)
    # List command
    list_parser = subparsers.add_parser(
        'ls', 
        help='Show available models for an API'
    )
    list_parser.add_argument(
        'api', 
        choices=SUPPORTED_APIS,
        help='API provider name'
    )
    list_parser.set_defaults(func=list_command)
    # Init command
    init_parser = subparsers.add_parser(
        'init', 
        help='Initialize environment variables'
    )
    init_parser.set_defaults(func=init_command)
    return parser

def main() -> None:
    parser = create_parser()
    args = parser.parse_args()
    if not hasattr(args, 'func'):
        parser.print_help()
        return
    try:
        args.func(args)
    except KeyboardInterrupt:
        rprint('\n[bold green]<<<[/bold green] [bold magenta]Bye![/bold magenta]')
    except Exception as e:
        rprint(f'[bold red]{ERROR_MESSAGES["unexpected_error"].format(error=e)}[/]')
        sys.exit(1)

if __name__ == '__main__':
    main()
